# """
#     执行该脚本前先准备以下内容

#     1.下载 https://github.com/kendryte/nncase/files/7594453/libomp140.x86_64.zip 并解压到 C:\Windows\System32
#     2.安装 dotnet-7 并添加环境变量
#     3.安装 python 3.10.11
#     4.创建 python 3.10.11 的虚拟环境
#     ----------------------------------------------
#         python -m venv venv
#         source venv/bin/activate     # Linux/MacOS
#         venv\Scripts\activate         # Windows
#         deactivate                    # 退出虚拟环境（可选）
#     ----------------------------------------------
#     5.进入 yolo 目录
#     6.安装 pip install ultralytics (不用开梯子)
#     7.开梯子
#      8. 把yolo香炉考下来
#     9.把脚本放在根目录运行 .py
#     10.需要提供数据集的路径
#     注意：如果报错找不到 datasets 请修改 C:\Users\<user_name>\AppData\Roaming\Ultralytics\settings.json 的 datasets_dir 为 yolo 目录，例如下面
#         "datasets_dir": "E:\\Project\\Github\\yolo_test\\yolo",

#     TODO:
#         1.后面考虑兼容 Linux/MacOS

# 最后：出现 project/train_output/weights/best.kmodel 就算成功了！！！

# """

import os
import sys
import subprocess
import requests
import zipfile
import shutil

def download_file(url, destination):
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    }
    if os.path.exists(destination):
        print(f"文件已存在，跳过下载：{destination}")
        return True

    print(f"下载文件 {url}...")
    try:
        response = requests.get(url, headers=headers, stream=True, allow_redirects=True)
        response.raise_for_status()
        os.makedirs(os.path.dirname(destination), exist_ok=True)
        with open(destination, "wb") as f:
            shutil.copyfileobj(response.raw, f)
        print(f"文件已下载到 {destination}！")
        return True
    except requests.exceptions.RequestException as e:
        print(f"下载失败：{e}")
        print(f"无法下载文件 {url}，请检查链接的合法性或稍后重试。")
        return False

def is_directory_empty(directory):
    """
    检查目录是否为空。
    如果目录不存在或为空，则返回 True；否则返回 False。
    """
    if not os.path.exists(directory):
        return True
    return len(os.listdir(directory)) == 0


def extract_zip(zip_path, extract_to, directory=None):
    """
    从 ZIP 文件中解压指定目录下的所有内容到目标目录。
    如果不指定目录，则解压整个 ZIP 文件。

    参数:
        zip_path (str): ZIP 文件的路径。
        directory (str, optional): ZIP 文件中需要解压的目录路径。默认为 None，表示解压整个 ZIP 文件。
        extract_to (str): 目标目录路径。默认为当前目录。
    """

    if not os.path.exists(zip_path):
        print(f"文件不存在，跳过解压：{zip_path}")
        return

    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        # 获取 ZIP 文件中的所有文件和目录
        all_files = zip_ref.namelist()

        if directory is None:
            # 解压整个 ZIP 文件
            print(f"解压整个 ZIP 文件到 {extract_to}...")
            zip_ref.extractall(extract_to)
        else:
            # 过滤出指定目录下的所有文件和子目录
            directory_files = [f for f in all_files if f.startswith(directory + os.sep)]

            if not directory_files:
                print(f"指定的目录 {directory} 在 ZIP 文件中不存在或为空。")
                return

            # 解压指定目录下的所有文件和子目录
            print(f"解压目录 {directory} 下的所有内容到 {extract_to}...")
            for file in directory_files:
                zip_ref.extract(file, extract_to)

        print(f"解压完成！")

def run_command(command):
    """
    运行命令
    """
    try:
        print(f"运行命令: {' '.join(command)}")
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        print(f"命令执行失败：{e}")


def train_yolo_model(data_config, model_file, epochs, imgsz, project_dir, experiment_name):
    command = [
        "yolo", "detect", "train",
        f"data={data_config}",
        f"model={model_file}",
        f"epochs={epochs}",
        f"imgsz={imgsz}",
        f"device=0",
        f"batch=-1",
        f"project={project_dir}",  # 正确的参数格式
        f"name={experiment_name}"   # 正确的参数格式
    ]
    run_command(command)

def export_yolo_model(runs_path):
    print("导出 YOLO 模型为 ONNX 格式...")
    model_path = os.path.join(runs_path, "weights", "best.pt")
    command = [
        "yolo", "export",
        f"model={model_path}",
        "format=onnx",
        "imgsz=320"
    ]
    run_command(command)
    print("ONNX 模型导出完成！")


def convert_to_kmodel(onnx_model_path, kmodel_output_dir, dataset_path):
    print("将 ONNX 模型转换为 Kmodel 格式...")
    os.makedirs(kmodel_output_dir, exist_ok=True)
    kmodel_output_path = os.path.join(kmodel_output_dir, "best.kmodel") # 指定输出的 kmodel 文件名

    command = [
        sys.executable,     # 使用当前 Python 解释器
        os.path.join("tools", "test_yolo11", "detect", "to_kmodel.py"),
        "--target", "k230",
        "--model", onnx_model_path,
        "--dataset", dataset_path,
        "--input_width", "320",
        "--input_height", "320",
        "--ptq_option", "0",
      
    ]
    run_command(command)
    print(f"Kmodel 转换完成！可能保存在: {os.path.dirname(kmodel_output_path)}")

    # 移动生成的 .kmodel 文件到指定的输出目录
    generated_kmodel_path = os.path.splitext(onnx_model_path)[0] + ".kmodel"
    final_kmodel_path = os.path.join(kmodel_output_dir, "best.kmodel")
    if os.path.exists(generated_kmodel_path):
        shutil.move(generated_kmodel_path, final_kmodel_path)
        print(f"Kmodel 文件已移动到: {final_kmodel_path}")
    else:
        print(f"警告：未找到生成的 Kmodel 文件: {generated_kmodel_path}")


def install_dependencies(requirements):
    """
    安装依赖库，如果依赖已安装则跳过。
    """
    print("检查并安装依赖库...")
    for dep in requirements:
        if os.path.isfile(dep):     # 检查是否为本地文件
            if not os.path.exists(dep):
                print(f"文件不存在，跳过安装：{dep}")
                continue
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", dep])
        except subprocess.CalledProcessError as e:
            print(f"安装依赖时发生错误：{e}")
    print("依赖安装完成！")


if __name__ == "__main__":
    # 定义根路径
    root_yolo_path = os.path.dirname(os.path.abspath(__file__))

    # --- 用户自定义数据集根路径 ---
    custom_dataset_root = "F:\\Code\\database\\tree"  # <--- 在这里修改你的数据集根路径
    # --- 用户自定义数据集根路径结束 ---

    downloads_path = os.path.join(root_yolo_path, "downloads")
    default_dataset_path = os.path.join(root_yolo_path, "datasets") # 原始脚本中的数据集路径
    tools_path = os.path.join(root_yolo_path, "tools")
    models_path = os.path.join(root_yolo_path, "models")

    # --- 自定义 Kmodel 输出目录 ---
    kmodel_output_base_dir = os.path.join(root_yolo_path, "project")
    kmodel_output_subdir = "test"
    kmodel_output_dir = os.path.join(kmodel_output_base_dir, kmodel_output_subdir)
    # --- 自定义 Kmodel 输出目录结束 ---

    # 创建必要的目录
    os.makedirs(downloads_path, exist_ok=True)
    os.makedirs(tools_path, exist_ok=True)
    os.makedirs(kmodel_output_dir, exist_ok=True) # 确保 Kmodel 输出目录存在

    # 下载文件
    download_urls = [
        ("https://github.com/kendryte/nncase/releases/download/v2.9.0/nncase_kpu-2.9.0-py2.py3-none-win_amd64.whl",
            os.path.join(downloads_path, "nncase_kpu-2.9.0-py2.py3-none-win_amd64.whl")),
        ("https://kendryte-download.canaan-creative.com/developer/k230/yolo_files/datasets.zip",
            os.path.join(downloads_path, "datasets.zip")),
        ("https://kendryte-download.canaan-creative.com/developer/k230/yolo_files/test_yolo11.zip",
            os.path.join(downloads_path, "test_yolo11.zip")),
        ("https://kendryte-download.canaan-creative.com/developer/k230/yolo_files/test_yolov5.zip",
            os.path.join(downloads_path, "test_yolov5.zip")),
        ("https://kendryte-download.canaan-creative.com/developer/k230/yolo_files/test_yolov8.zip",
            os.path.join(downloads_path, "test_yolov8.zip"))
    ]
    failed_downloads = []
    for url, destination in download_urls:
        if not download_file(url, destination):
            failed_downloads.append((url, destination))
    if failed_downloads:
        print("\n以下文件下载失败：")
        for url, destination in failed_downloads:
            print(f"下载地址：{url}")
            print(f"目标位置：{destination}")
        print("请检查网络连接或离线下载到指定位置！！！\n再次重试！！！")
        sys.exit(1)

    # 解压文件
    zip_files = [
        (os.path.join(downloads_path, "datasets.zip"), root_yolo_path, None),
        (os.path.join(downloads_path, "test_yolo11.zip"), tools_path, None)
    ]
    for zip_path, extract_to, directory in zip_files:
        extract_zip(zip_path, extract_to, directory)

    # 安装依赖
    requirements = [
        "onnx",
        "onnxruntime",
        "onnxsim",
        "nncase==2.9.0",
        os.path.join(downloads_path, "nncase_kpu-2.9.0-py2.py3-none-win_amd64.whl")
    ]
    install_dependencies(requirements)

    # 定义数据集路径
    data_config_path = os.path.join(custom_dataset_root, "num_yolo.yaml")
    train_images_path = os.path.join(custom_dataset_root,"num_yolo", "images", "train")
    val_images_path = os.path.join(custom_dataset_root, "num_yolo","images", "val")
    test_images_path = os.path.join(custom_dataset_root, "num_yolo","images", "test")

    # 确保数据集配置文件存在
    if not os.path.exists(data_config_path):
        print(f"错误：数据集配置文件 {data_config_path} 不存在，请创建并配置它。")
        sys.exit(1)

    # 定义训练输出路径
    train_output_project = os.path.join(root_yolo_path, "project")
    train_output_name = "train_output"

    # 训练模型
    train_yolo_model(
        data_config=data_config_path,
        model_file=os.path.join(models_path, "yolo11n.pt"),
        epochs=2,  # 减少 epochs 以加快测试
        imgsz=320,
        project_dir=train_output_project,
        experiment_name=train_output_name
    )

    # 导出 ONNX 模型
    runs_path = os.path.join(train_output_project, train_output_name)
    export_yolo_model(runs_path)

    # 转换为 Kmodel
    onnx_model_path = os.path.join(runs_path, "weights", "best.onnx")
    convert_to_kmodel(onnx_model_path, kmodel_output_dir, test_images_path)

    print("\n脚本执行完毕！Kmodel 文件已生成在:", os.path.join(kmodel_output_dir, "best.kmodel"))
    print(f"YOLO 训练结果保存在: {runs_path}")
    print("你可以通过修改脚本中的 'custom_dataset_root' 变量来更改数据集路径。")